#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import zipfile
import tarfile
import gzip
import shutil
import py7zr
import threading
from typing import Callable, List, Optional, Tuple, Union

class CompressionHandler:
    """处理各种压缩和解压缩操作的类"""
    
    # 支持的压缩格式
    SUPPORTED_FORMATS = {
        'zip': '.zip',
        '7z': '.7z',
        'tar': '.tar',
        'tar.gz': '.tar.gz',
        'tgz': '.tgz',
        'gz': '.gz'
    }
    
    @staticmethod
    def get_supported_formats() -> dict:
        """获取支持的压缩格式"""
        return CompressionHandler.SUPPORTED_FORMATS
    
    @staticmethod
    def get_format_from_filename(filename: str) -> Optional[str]:
        """从文件名中获取压缩格式"""
        for format_name, extension in CompressionHandler.SUPPORTED_FORMATS.items():
            if filename.endswith(extension):
                return format_name
        return None
    
    @staticmethod
    def compress_files(
        files: List[str], 
        output_path: str, 
        format_name: str, 
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """
        压缩文件或文件夹
        
        Args:
            files: 要压缩的文件或文件夹路径列表
            output_path: 输出压缩文件的路径
            format_name: 压缩格式名称
            progress_callback: 进度回调函数，接收(current, total)参数
            
        Returns:
            (成功标志, 错误信息)
        """
        try:
            # 确保输出目录存在
            output_dir = os.path.dirname(output_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)
            
            # 根据格式选择压缩方法
            if format_name == 'zip':
                return CompressionHandler._compress_zip(files, output_path, progress_callback)
            elif format_name == '7z':
                return CompressionHandler._compress_7z(files, output_path, progress_callback)
            elif format_name in ['tar', 'tar.gz', 'tgz']:
                return CompressionHandler._compress_tar(files, output_path, format_name, progress_callback)
            elif format_name == 'gz':
                if len(files) != 1:
                    return False, "gz格式只支持压缩单个文件"
                return CompressionHandler._compress_gz(files[0], output_path, progress_callback)
            else:
                return False, f"不支持的压缩格式: {format_name}"
        except Exception as e:
            return False, f"压缩过程中发生错误: {str(e)}"
    
    @staticmethod
    def _compress_zip(
        files: List[str], 
        output_path: str, 
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """ZIP格式压缩实现"""
        try:
            # 计算总文件数量用于进度显示
            total_files = 0
            for file_path in files:
                if os.path.isfile(file_path):
                    total_files += 1
                else:
                    for root, _, filenames in os.walk(file_path):
                        total_files += len(filenames)
            
            processed_files = 0
            
            with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for file_path in files:
                    if os.path.isfile(file_path):
                        # 添加单个文件
                        zipf.write(file_path, os.path.basename(file_path))
                        processed_files += 1
                        if progress_callback:
                            progress_callback(processed_files, total_files)
                    else:
                        # 添加整个目录
                        base_path = os.path.dirname(file_path)
                        for root, _, filenames in os.walk(file_path):
                            for filename in filenames:
                                file_full_path = os.path.join(root, filename)
                                # 计算相对路径，保持目录结构
                                if base_path:
                                    arcname = os.path.relpath(file_full_path, base_path)
                                else:
                                    arcname = file_full_path
                                zipf.write(file_full_path, arcname)
                                processed_files += 1
                                if progress_callback:
                                    progress_callback(processed_files, total_files)
            
            return True, ""
        except Exception as e:
            return False, f"ZIP压缩错误: {str(e)}"
    
    @staticmethod
    def _compress_7z(
        files: List[str], 
        output_path: str, 
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """7Z格式压缩实现"""
        try:
            # 计算总文件数量
            total_files = 0
            file_list = []
            
            for file_path in files:
                if os.path.isfile(file_path):
                    total_files += 1
                    file_list.append((file_path, os.path.basename(file_path)))
                else:
                    base_path = os.path.dirname(file_path)
                    for root, _, filenames in os.walk(file_path):
                        for filename in filenames:
                            file_full_path = os.path.join(root, filename)
                            if base_path:
                                arcname = os.path.relpath(file_full_path, base_path)
                            else:
                                arcname = file_full_path
                            total_files += 1
                            file_list.append((file_full_path, arcname))
            
            # 创建7z文件
            with py7zr.SevenZipFile(output_path, mode='w') as archive:
                for i, (file_path, arcname) in enumerate(file_list):
                    archive.write(file_path, arcname)
                    if progress_callback:
                        progress_callback(i + 1, total_files)
            
            return True, ""
        except Exception as e:
            return False, f"7Z压缩错误: {str(e)}"
    
    @staticmethod
    def _compress_tar(
        files: List[str], 
        output_path: str, 
        format_name: str,
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """TAR/TAR.GZ格式压缩实现"""
        try:
            # 确定压缩模式
            mode = 'w'
            if format_name in ['tar.gz', 'tgz']:
                mode = 'w:gz'
            
            # 计算总文件数量
            total_files = 0
            for file_path in files:
                if os.path.isfile(file_path):
                    total_files += 1
                else:
                    for root, _, filenames in os.walk(file_path):
                        total_files += len(filenames)
            
            processed_files = 0
            
            with tarfile.open(output_path, mode) as tar:
                for file_path in files:
                    if os.path.isfile(file_path):
                        # 添加单个文件
                        tar.add(file_path, arcname=os.path.basename(file_path))
                        processed_files += 1
                        if progress_callback:
                            progress_callback(processed_files, total_files)
                    else:
                        # 添加整个目录
                        base_path = os.path.dirname(file_path)
                        for root, _, filenames in os.walk(file_path):
                            for filename in filenames:
                                file_full_path = os.path.join(root, filename)
                                if base_path:
                                    arcname = os.path.relpath(file_full_path, base_path)
                                else:
                                    arcname = file_full_path
                                tar.add(file_full_path, arcname=arcname)
                                processed_files += 1
                                if progress_callback:
                                    progress_callback(processed_files, total_files)
            
            return True, ""
        except Exception as e:
            return False, f"TAR压缩错误: {str(e)}"
    
    @staticmethod
    def _compress_gz(
        file_path: str, 
        output_path: str, 
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """GZ格式压缩实现（单文件）"""
        try:
            # 获取文件大小用于进度显示
            file_size = os.path.getsize(file_path)
            
            # 分块读取以处理大文件
            chunk_size = 1024 * 1024  # 1MB
            processed_size = 0
            
            with open(file_path, 'rb') as f_in:
                with gzip.open(output_path, 'wb') as f_out:
                    while True:
                        chunk = f_in.read(chunk_size)
                        if not chunk:
                            break
                        f_out.write(chunk)
                        processed_size += len(chunk)
                        if progress_callback:
                            progress_callback(processed_size, file_size)
            
            return True, ""
        except Exception as e:
            return False, f"GZ压缩错误: {str(e)}"
    
    @staticmethod
    def extract_archive(
        archive_path: str, 
        output_dir: str, 
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """
        解压缩文件
        
        Args:
            archive_path: 压缩文件路径
            output_dir: 解压输出目录
            progress_callback: 进度回调函数，接收(current, total)参数
            
        Returns:
            (成功标志, 错误信息)
        """
        try:
            # 确保输出目录存在
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            
            # 获取文件格式
            format_name = CompressionHandler.get_format_from_filename(archive_path)
            if not format_name:
                return False, "无法识别的压缩文件格式"
            
            # 根据格式选择解压方法
            if format_name == 'zip':
                return CompressionHandler._extract_zip(archive_path, output_dir, progress_callback)
            elif format_name == '7z':
                return CompressionHandler._extract_7z(archive_path, output_dir, progress_callback)
            elif format_name in ['tar', 'tar.gz', 'tgz']:
                return CompressionHandler._extract_tar(archive_path, output_dir, progress_callback)
            elif format_name == 'gz':
                return CompressionHandler._extract_gz(archive_path, output_dir, progress_callback)
            else:
                return False, f"不支持的压缩格式: {format_name}"
        except Exception as e:
            return False, f"解压过程中发生错误: {str(e)}"
    
    @staticmethod
    def _extract_zip(
        archive_path: str, 
        output_dir: str, 
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """ZIP格式解压实现"""
        try:
            with zipfile.ZipFile(archive_path, 'r') as zipf:
                # 获取文件列表用于进度显示
                file_list = zipf.namelist()
                total_files = len(file_list)
                
                # 解压所有文件
                for i, file in enumerate(file_list):
                    zipf.extract(file, output_dir)
                    if progress_callback:
                        progress_callback(i + 1, total_files)
            
            return True, ""
        except zipfile.BadZipFile:
            return False, "无效的ZIP文件"
        except Exception as e:
            return False, f"ZIP解压错误: {str(e)}"
    
    @staticmethod
    def _extract_7z(
        archive_path: str, 
        output_dir: str, 
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """7Z格式解压实现"""
        try:
            with py7zr.SevenZipFile(archive_path, mode='r') as archive:
                # 获取文件列表
                file_list = archive.getnames()
                total_files = len(file_list)
                
                # 使用回调进行解压
                def progress_monitor(total, completed):
                    if progress_callback:
                        progress_callback(completed, total)
                
                # 解压所有文件
                archive.extractall(path=output_dir, progress_callback=progress_monitor)
            
            return True, ""
        except Exception as e:
            return False, f"7Z解压错误: {str(e)}"
    
    @staticmethod
    def _extract_tar(
        archive_path: str, 
        output_dir: str, 
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """TAR/TAR.GZ格式解压实现"""
        try:
            with tarfile.open(archive_path, 'r:*') as tar:
                # 获取文件列表用于进度显示
                file_list = tar.getnames()
                total_files = len(file_list)
                
                # 解压所有文件
                for i, file in enumerate(file_list):
                    tar.extract(file, output_dir)
                    if progress_callback:
                        progress_callback(i + 1, total_files)
            
            return True, ""
        except tarfile.ReadError:
            return False, "无效的TAR文件"
        except Exception as e:
            return False, f"TAR解压错误: {str(e)}"
    
    @staticmethod
    def _extract_gz(
        archive_path: str, 
        output_dir: str, 
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Tuple[bool, str]:
        """GZ格式解压实现（单文件）"""
        try:
            # 确定输出文件名（去掉.gz扩展名）
            output_filename = os.path.basename(archive_path)
            if output_filename.endswith('.gz'):
                output_filename = output_filename[:-3]
            
            output_path = os.path.join(output_dir, output_filename)
            
            # 获取文件大小用于进度显示
            file_size = os.path.getsize(archive_path)
            
            # 分块读取以处理大文件
            chunk_size = 1024 * 1024  # 1MB
            processed_size = 0
            
            with gzip.open(archive_path, 'rb') as f_in:
                with open(output_path, 'wb') as f_out:
                    while True:
                        chunk = f_in.read(chunk_size)
                        if not chunk:
                            break
                        f_out.write(chunk)
                        processed_size += len(chunk)
                        if progress_callback:
                            progress_callback(processed_size, file_size)
            
            return True, ""
        except Exception as e:
            return False, f"GZ解压错误: {str(e)}"
    
    @staticmethod
    def run_in_thread(
        func: Callable, 
        args: tuple, 
        on_complete: Optional[Callable[[bool, str], None]] = None
    ) -> threading.Thread:
        """
        在单独的线程中运行函数，避免阻塞GUI
        
        Args:
            func: 要运行的函数
            args: 函数参数
            on_complete: 完成回调，接收(success, error_message)参数
            
        Returns:
            线程对象
        """
        def thread_func():
            result = func(*args)
            if on_complete:
                on_complete(*result)
        
        thread = threading.Thread(target=thread_func)
        thread.daemon = True
        thread.start()
        return thread